import os
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

from FlagEmbedding import BGEM3FlagModel
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
import torch

# 加载BGE-M3模型
model = BGEM3FlagModel('BAAI/bge-m3', use_fp16=True)

# 示例文本数据
corpus = [
    "自然语言处理是人工智能的一个重要领域",
    "大型语言模型在许多NLP任务中表现出色",
    "计算机视觉处理图像和视频数据",
    "推荐系统基于用户行为预测偏好",
    "信息检索关注如何高效获取相关信息"
]

# 查询文本
query = "什么技术用于文本分析？"

# 生成文本向量（密集检索）
def get_dense_embeddings(texts):
    embeddings = model.encode(texts, batch_size=128, max_length=8192)['dense_vecs']
    return embeddings

# 生成稀疏向量（稀疏检索）
def get_sparse_embeddings(texts):
    output = model.encode(texts, return_dense=True, return_sparse=True, return_colbert_vecs=False)
    return output['lexical_weights']

# 密集检索示例
def dense_retrieval(query, corpus):
    # 生成查询和语料库的密集向量
    query_embedding = get_dense_embeddings([query])
    corpus_embeddings = get_dense_embeddings(corpus)
    
    # 计算余弦相似度
    similarities = cosine_similarity(query_embedding, corpus_embeddings)[0]
    
    # 按相似度排序
    results = sorted(zip(corpus, similarities), key=lambda x: x[1], reverse=True)
    
    return results

# 稀疏检索示例
def sparse_retrieval(query, corpus):
    # 生成查询和语料库的稀疏向量
    query_sparse = get_sparse_embeddings([query])[0]
    corpus_sparse = get_sparse_embeddings(corpus)
    
    # 计算词法匹配分数
    scores = []
    for doc_sparse in corpus_sparse:
        score = model.compute_lexical_matching_score(query_sparse, doc_sparse)
        scores.append(score)
    
    # 按分数排序
    results = sorted(zip(corpus, scores), key=lambda x: x[1], reverse=True)
    
    return results

# 混合检索示例（结合密集和稀疏检索）
def hybrid_retrieval(query, corpus, alpha=0.7):
    # 密集检索结果
    dense_results = dense_retrieval(query, corpus)
    dense_scores = {doc: score for doc, score in dense_results}
    
    # 稀疏检索结果
    sparse_results = sparse_retrieval(query, corpus)
    sparse_scores = {doc: score for doc, score in sparse_results}
    
    # 混合分数
    hybrid_scores = []
    for doc in corpus:
        hybrid_score = alpha * dense_scores[doc] + (1 - alpha) * sparse_scores[doc]
        hybrid_scores.append((doc, hybrid_score))
    
    # 按混合分数排序
    results = sorted(hybrid_scores, key=lambda x: x[1], reverse=True)
    
    return results

# 执行检索
print("===== 密集检索结果 =====")
for doc, score in dense_retrieval(query, corpus):
    print(f"相似度: {score:.4f} | 文档: {doc}")

print("\n===== 稀疏检索结果 =====")
for doc, score in sparse_retrieval(query, corpus):
    print(f"相似度: {score:.4f} | 文档: {doc}")

print("\n===== 混合检索结果 =====")
for doc, score in hybrid_retrieval(query, corpus):
    print(f"相似度: {score:.4f} | 文档: {doc}")